from flask import Blueprint, request, jsonify, Response, current_app
from openai import OpenAI
from config import Config
from .models import db, Conversation, User  # 导入User模型
import time
import json
import logging
from datetime import datetime
from flask_jwt_extended import jwt_required, get_jwt_identity, get_jwt, verify_jwt_in_request  # 导入JWT相关函数

# 配置日志
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# 创建一个装饰器函数来捕获和记录JWT错误
def jwt_error_handler(fn):
    def wrapper(*args, **kwargs):
        try:
            # 尝试验证JWT
            auth_header = request.headers.get('Authorization', 'No Authorization header')
            logger.debug(f"Auth header: {auth_header[:20]}...")
            
            verify_jwt_in_request()
            current_user_id = get_jwt_identity()
            logger.debug(f"JWT verified for user ID: {current_user_id}")
            
            return fn(*args, **kwargs)
        except Exception as e:
            logger.error(f"JWT authentication error: {str(e)}")
            logger.error(f"Headers: {dict(request.headers)}")
            return jsonify({'error': f'认证失败: {str(e)}'}), 401
    wrapper.__name__ = fn.__name__
    return wrapper

main = Blueprint('main', __name__)

def interleave_messages(messages):
    if not messages:
        return []
    result = [messages[0]]
    for msg in messages[1:]:
        if msg['role'] == result[-1]['role']:
            continue  # 跳过连续相同角色
        result.append(msg)
    return result

@main.route('/api/chat', methods=['POST'])
@jwt_error_handler
def chat():
    current_user_id = get_jwt_identity()  # 获取当前登录用户的ID（字符串）
    user = User.query.get(int(current_user_id))  # 转换为整数查询
    if not user:
        return jsonify({'error': '用户不存在'}), 401

    try:
        data = request.json
        messages = data.get('messages', [])
        messages = interleave_messages(messages)
        model = data.get('model', 'deepseek-chat')
        conversation_id = data.get('conversation_id')  # 现在这是整数 id 而非 uid

        # 验证消息格式
        if not isinstance(messages, list):
            return jsonify({'error': '消息格式错误，必须是列表'}), 400

        # 验证每条消息
        for i, msg in enumerate(messages):
            if not isinstance(msg, dict):
                return jsonify({'error': f'消息 #{i+1} 格式错误，必须是对象'}), 400
            if 'role' not in msg or 'content' not in msg:
                return jsonify({'error': f'消息 #{i+1} 缺少必须的 role 或 content 字段'}), 400
            if msg['role'] not in ['user', 'assistant', 'system']:
                return jsonify({'error': f'消息 #{i+1} 的角色无效，必须是 user, assistant 或 system'}), 400
            if not isinstance(msg['content'], str):
                return jsonify({'error': f'消息 #{i+1} 的内容必须是字符串'}), 400

        cleaned_messages = [{'role': msg['role'], 'content': msg['content']} for msg in messages]

        if not cleaned_messages:
            return jsonify({'error': '消息不能为空'}), 400

        client = OpenAI(api_key=Config.OPENAI_API_KEY, base_url="https://api.deepseek.com/v1")
        app = current_app._get_current_object()

        def generate():
            try:
                response = client.chat.completions.create(
                    model=model,
                    messages=cleaned_messages,
                    temperature=0.7,
                    stream=True,
                    max_tokens=4096
                )
                full_assistant_content = ""
                full_reasoning_content = ""
                for chunk in response:
                    try:
                        if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                            content = chunk.choices[0].delta.reasoning_content
                            full_reasoning_content += content
                            yield json.dumps({
                                'type': 'reasoning',
                                'content': content
                            }) + '\n'

                        if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                            content = chunk.choices[0].delta.content
                            full_assistant_content += content
                            yield json.dumps({
                                'type': 'content',
                                'content': content
                            }) + '\n'

                    except Exception as chunk_error:
                        logger.error(f"Error processing chunk: {chunk_error}")
                        yield json.dumps({'error': str(chunk_error)}) + '\n'

                try:
                    with app.app_context():
                        assistant_reply = {'role': 'assistant', 'content': full_assistant_content}
                        if full_reasoning_content:
                            assistant_reply['reasoning_content'] = full_reasoning_content

                        conversation_messages = cleaned_messages + [assistant_reply]

                        conversation = None
                        # 如果提供了会话ID，尝试获取现有会话
                        if conversation_id:
                            conversation = Conversation.query.filter_by(id=conversation_id, user_id=int(current_user_id)).first()
                            if conversation:
                                conversation.messages = conversation_messages
                                conversation.updated_at = datetime.utcnow()
                            else:
                                logger.warning(f"Conversation {conversation_id} not found or doesn't belong to user {current_user_id}. Creating new.")
                                conversation = Conversation.create_conversation(int(current_user_id), conversation_messages)
                        else:
                            conversation = Conversation.create_conversation(int(current_user_id), conversation_messages)

                        db.session.add(conversation)
                        db.session.commit()

                        yield json.dumps({
                            'type': 'conversation_id',
                            'conversation_id': conversation.id
                        }) + '\n'

                except Exception as db_error:
                    logger.error(f"Error saving conversation for user {current_user_id}: {db_error}")
                    yield json.dumps({'error': f"保存对话失败: {str(db_error)}"}) + '\n'

            except Exception as e:
                logger.error(f"Error in generate for user {current_user_id}: {e}")
                yield json.dumps({'error': str(e)}) + '\n'

        return Response(generate(), mimetype='text/event-stream')

    except Exception as e:
        logger.error(f"Error in chat endpoint for user {current_user_id}: {e}")
        return jsonify({'error': str(e)}), 500

@main.route('/api/conversations', methods=['GET'])
@jwt_required()
def get_conversations():
    """获取当前用户的所有对话历史"""
    current_user_id = get_jwt_identity()  # 字符串
    try:
        page = request.args.get('page', 1, type=int)
        per_page = request.args.get('per_page', 10, type=int)

        pagination = Conversation.query.filter_by(user_id=int(current_user_id)).order_by(
            Conversation.updated_at.desc()
        ).paginate(
            page=page, per_page=per_page, error_out=False
        )

        conversations = [conv.to_dict(include_messages=False) for conv in pagination.items]
        
        # 如果没有历史对话，返回一个欢迎信息
        if not conversations and page == 1:
            return jsonify({
                'conversations': [],
                'total': 0,
                'pages': 0,
                'current_page': 1,
                'welcome_message': '欢迎使用 AI 问答助手！'  # 添加欢迎信息
            })

        return jsonify({
            'conversations': conversations,
            'total': pagination.total,
            'pages': pagination.pages,
            'current_page': page
        })
    except Exception as e:
        logger.error(f"Error getting conversations for user {current_user_id}: {e}")
        return jsonify({'error': str(e)}), 500

@main.route('/api/conversations/<int:conversation_id>', methods=['GET'])
@jwt_required()  # 需要登录
def get_conversation(conversation_id):
    """获取特定对话的历史记录"""
    current_user_id = get_jwt_identity()  # 现在这会返回字符串
    try:
        conversation = Conversation.query.filter_by(id=conversation_id, user_id=int(current_user_id)).first()
        if conversation:
            return jsonify(conversation.to_dict(include_messages=True))
        return jsonify({'error': '对话不存在或无权访问'}), 404
    except Exception as e:
        logger.error(f"Error getting conversation {conversation_id} for user {current_user_id}: {e}")
        return jsonify({'error': str(e)}), 500

@main.route('/api/conversations/<int:conversation_id>', methods=['DELETE'])
@jwt_required()  # 需要登录
def delete_conversation(conversation_id):
    """删除特定对话"""
    current_user_id = get_jwt_identity()  # 现在这会返回字符串
    try:
        conversation = Conversation.query.filter_by(id=conversation_id, user_id=int(current_user_id)).first()
        if conversation:
            db.session.delete(conversation)
            db.session.commit()
            logger.info(f"Conversation {conversation_id} deleted by user {current_user_id}.")
            return jsonify({'message': '对话已删除'})
        return jsonify({'error': '对话不存在或无权访问'}), 404
    except Exception as e:
        logger.error(f"Error deleting conversation {conversation_id} for user {current_user_id}: {e}")
        return jsonify({'error': str(e)}), 500

@main.route('/api/conversations', methods=['DELETE'])
@jwt_required()  # 需要登录
def clear_conversations():
    """清空当前用户的所有对话"""
    current_user_id = get_jwt_identity()  # 现在这会返回字符串
    try:
        num_deleted = Conversation.query.filter_by(user_id=int(current_user_id)).delete()
        db.session.commit()
        logger.info(f"{num_deleted} conversations deleted for user {current_user_id}.")
        return jsonify({'message': f'{num_deleted} 个对话已清空'})
    except Exception as e:
        logger.error(f"Error clearing conversations for user {current_user_id}: {e}")
        return jsonify({'error': str(e)}), 500

@main.route('/api/models', methods=['GET'])
def get_models():
    """获取可用模型列表"""
    models = [
        {
            'id': 'deepseek-chat',
            'name': 'DeepSeek Chat',
            'description': 'DeepSeek 通用对话模型'
        },
        {
            'id': 'deepseek-reasoner',
            'name': 'DeepSeek Reasoner',
            'description': 'DeepSeek 推理增强模型'
        }
    ]
    return jsonify(models)

@main.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy'})

@main.route('/api/debug/auth-test', methods=['GET'])
def auth_test():
    """测试端点，检查认证信息"""
    auth_header = request.headers.get('Authorization', 'No Authorization header')
    
    # 基本请求信息
    info = {
        'auth_header': auth_header,
        'auth_header_valid': False,
        'ip': request.remote_addr,
        'method': request.method,
        'path': request.path,
        'time': datetime.now().isoformat()
    }
    
    # 检查格式
    if auth_header and auth_header.startswith('Bearer '):
        token = auth_header[7:].strip()
        info['auth_header_valid'] = len(token) > 0
        info['token_length'] = len(token) if len(token) > 0 else 0
        
        if len(token) > 0:
            info['token_preview'] = f"{token[:10]}...{token[-10:]}" if len(token) > 20 else token
            
            # 尝试解析但不验证 JWT
            try:
                from flask_jwt_extended import decode_token
                decoded = decode_token(token, allow_expired=True)
                info['token_decode_success'] = True
                info['token_sub'] = decoded.get('sub')
                info['token_exp'] = decoded.get('exp')
                info['token_iat'] = decoded.get('iat')
            except Exception as e:
                info['token_decode_success'] = False
                info['token_decode_error'] = str(e)
    
    return jsonify(info)